package com.bw;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.dataproc.SplitBatchOp;
import com.alibaba.alink.operator.batch.source.CsvSourceBatchOp;
import com.alibaba.alink.params.feature.HasEncodeWithoutWoe;
import com.alibaba.alink.pipeline.dataproc.vector.VectorAssembler;
import com.alibaba.alink.pipeline.feature.OneHotEncoder;
import com.alibaba.alink.pipeline.regression.DecisionTreeRegressor;
import com.alibaba.alink.pipeline.regression.LinearRegression;
import com.alibaba.alink.pipeline.regression.LinearRegressionModel;

public class Test3 {
    public static void main(String[] args) throws Exception {
        BatchOperator.setParallelism(1);
        String schemaStr = "car_ID int,symboling int,CarName string,fueltype string,aspiration string,doornumber string,carbody string,drivewheel string,enginelocation string,wheelbase double,carlength double,carwidth double,carheight double,curbweight int,enginetype string,cylindernumber string,enginesize int,fuelsystem string,boreratio double,stroke double,compressionratio double,horsepower int,peakrpm int,citympg int,highwaympg int,price double";

        //CarName  fueltype  aspiration doornumber carbody drivewheel
        CsvSourceBatchOp source = new CsvSourceBatchOp()
                .setFilePath("datafile/test2.csv")
                .setSchemaStr(schemaStr)
                .setFieldDelimiter(",")
                .setIgnoreFirstLine(true);
//        source.print();

        //编码
        OneHotEncoder one_hot_06 = new OneHotEncoder()
                .setSelectedCols("CarName", "fueltype", "aspiration", "doornumber", "carbody", "drivewheel")
                .setOutputCols("vec");

        BatchOperator<?> source_06 = one_hot_06.fit(source).transform(source);
//        source_06.print();


        // 选择特征
        String[] features = new String[]{"symboling","vec","highwaympg"};
        // 目标列
        String lable = "price";


        // 向量聚合
        VectorAssembler res = new VectorAssembler()
                .setSelectedCols(features)
                .setOutputCol("vec1");
        BatchOperator<?> VectorResult =res.transform(source_06);
//        VectorResult.print();

        // 拆分数据集
        BatchOperator<?> trainData = new SplitBatchOp().setFraction(0.8).linkFrom(VectorResult);
        BatchOperator<?> testData = trainData.getSideOutput(0);

        // 线性回归
//        LinearRegression lr = new LinearRegression()
////                .setFeatureCols(features)
//                .setVectorCol("vec1")
//                .setLabelCol(lable)
//                .setPredictionCol("pred")
//                .enableLazyPrintModelInfo("ModelInfo");
//        // 训练模型
//        LinearRegressionModel lr_model = lr.fit(trainData);
//        BatchOperator<?> lr_result=lr_model.transform(testData);
//        lr_result.print();

        // 用决策树回归
        BatchOperator<?> dt_result = new DecisionTreeRegressor()
                .setPredictionCol("pred")
                .setLabelCol(lable)
                .setFeatureCols(features)
                .fit(trainData)
                .transform(testData);
        dt_result.print();


    }
}
